import pyaudio
import wave
import time
import sys
import numpy as np
import threading
import keyboard  # 需要额外安装: pip install keyboard

def list_audio_devices():
    """
    列出所有可用的音频输入设备
    :return: 可用的音频输入设备列表
    """
    audio = pyaudio.PyAudio()
    info = audio.get_host_api_info_by_index(0)
    num_devices = info.get('deviceCount')
    
    print("可用的音频输入设备:")
    input_devices = []
    
    for i in range(num_devices):
        device_info = audio.get_device_info_by_index(i)
        
        if device_info.get('maxInputChannels') > 0:  # 如果设备支持输入
            device_dict={i:device_info.get('name')}
            input_devices.append(device_dict)
            print(f"设备 ID {i}: {device_info.get('name')}")
    
    audio.terminate()
    
    if not input_devices:
        print("未检测到任何音频输入设备！")
        
    return input_devices

def find_audio_device(input_devices, keyword="audio"):
    """
    在输入设备列表中查找包含指定关键词的设备
    :param input_devices: 输入设备列表
    :param keyword: 关键词，默认为"audio"
    :return: 设备ID
    """
    device_idx = 0
    audio_device_found = False
    
    for i in input_devices:
        for k, v in i.items():
            if keyword.lower() in str(v).lower():  # 不区分大小写
                device_idx = k
                print(f"找到包含'{keyword}'的设备: ID {k}: {v}")
                audio_device_found = True
                break
        if audio_device_found:
            break
    
    # 如果没有找到包含关键词的设备，使用第一个可用设备
    if not audio_device_found:
        print(f"未找到包含'{keyword}'的设备，使用第一个可用设备")
        first_device = input_devices[0]
        for k, v in first_device.items():
            device_idx = k
            print(f"使用设备: ID {k}: {v}")
            break
            
    return device_idx

def record_microphone_to_wav(output_file='recording.wav', record_seconds=10, device_index=None):
    """
    录制麦克风音频流并保存为WAV文件（不支持静音检测-录制指定的时长）
    :param output_file: 输出文件名(.wav)
    :param record_seconds: 录制时长(秒)，0表示持续录制直到手动停止
    :param device_index: 音频设备索引，None表示使用默认设备
    """
    # 音频参数设置
    FORMAT = pyaudio.paInt16  # 16-bit采样
    CHANNELS = 1              # 单声道
    RATE = 44100              # 采样率44.1kHz
    CHUNK = 1024              # 每次读取的帧数
    
    audio = pyaudio.PyAudio()
    
    # 打开音频流
    stream = audio.open(format=FORMAT,
                        channels=CHANNELS,
                        rate=RATE,
                        input=True,
                        input_device_index=device_index,
                        frames_per_buffer=CHUNK)
    
    print("开始录制...")
    if record_seconds > 0:
        print(f"将自动录制 {record_seconds} 秒")
    else:
        print("持续录制中，按Ctrl+C停止")
    
    frames = []
    
    try:
        if record_seconds > 0:
            # 自动停止录制
            for i in range(0, int(RATE / CHUNK * record_seconds)):
                data = stream.read(CHUNK, exception_on_overflow=False)
                frames.append(data)
        else:
            # 手动停止录制
            while True:
                data = stream.read(CHUNK, exception_on_overflow=False)
                frames.append(data)
                # 可以添加实时处理逻辑
                
    except KeyboardInterrupt:
        print("\n录制被用户中断")
    except Exception as e:
        print(f"录制出错: {str(e)}")
    finally:
        # 停止并关闭流
        print("录制结束，保存文件中...")
        stream.stop_stream()
        stream.close()
        audio.terminate()
        
        # 保存为WAV文件
        with wave.open(output_file, 'wb') as wf:
            wf.setnchannels(CHANNELS)
            wf.setsampwidth(audio.get_sample_size(FORMAT))
            wf.setframerate(RATE)
            wf.writeframes(b''.join(frames))
        
        print(f"音频已保存为 {output_file}")

def record_with_silence_detection(output_file='recording.wav', max_record_seconds=30, device_index=None, silence_threshold=10, silence_duration=2, device_keyword="audio"):
    """
    录制麦克风音频流并保存为WAV文件，支持静音检测自动停止（有2秒声音则停止和最大录制时长）
    :param output_file: 输出文件名(.wav)
    :param max_record_seconds: 最大录制时长(秒)
    :param device_index: 音频设备索引，None表示自动选择
    :param silence_threshold: 静音阈值，值越小越灵敏
    :param silence_duration: 持续静音多少秒后自动停止，默认2秒
    :param device_keyword: 设备选择关键词，默认为"audio"
    :return: 是否检测到声音
    """
    # 如果没有指定设备索引，则自动选择设备
    if device_index is None:
        # 列出所有可用设备
        input_devices = list_audio_devices()
        
        if not input_devices:
            print("无法继续录音，请检查您的麦克风连接或权限设置。")
            return False
        
        # 查找包含关键词的设备
        device_index = find_audio_device(input_devices, device_keyword)
        print(f"将使用设备ID {device_index} 录制音频到文件 {output_file}")
    
    # 音频参数设置
    FORMAT = pyaudio.paInt16  # 16-bit采样
    CHANNELS = 1              # 单声道
    RATE = 44100              # 采样率44.1kHz
    CHUNK = 1024              # 每次读取的帧数
    
    audio = pyaudio.PyAudio()
    
    # 打开音频流
    stream = audio.open(format=FORMAT,
                        channels=CHANNELS,
                        rate=RATE,
                        input=True,
                        input_device_index=device_index,
                        frames_per_buffer=CHUNK)
    
    print("开始录制...")
    print(f"最大录制时长 {max_record_seconds} 秒")
    print(f"静音检测: {silence_duration}秒内没有声音将自动停止")
    
    frames = []
    silent_chunks = 0
    voice_detected = False
    max_silent_chunks = int(silence_duration * RATE / CHUNK)  # 2秒内的静音块数
    
    # 先收集一些背景噪音样本来校准
    print("校准中，请保持安静...")
    background_samples = []
    for i in range(0, int(RATE / CHUNK * 1)):  # 收集1秒的背景噪音
        data = stream.read(CHUNK, exception_on_overflow=False)
        background_samples.append(data)
    
    # 计算背景噪音水平
    background_noise = 0
    for sample in background_samples:
        audio_data = np.frombuffer(sample, dtype=np.int16)
        background_noise += np.abs(audio_data).mean()
    
    background_noise = background_noise / len(background_samples) if background_samples else 0
    # 将阈值设为背景噪音的2倍，但不小于预设值
    adjusted_threshold = max(silence_threshold, background_noise * 2)
    
    print(f"背景噪音水平: {background_noise:.2f}")
    print(f"调整后的静音阈值: {adjusted_threshold:.2f}")
    
    # 初始静音计数，如果一开始就是静音，2秒后就会自动结束
    silent_chunks = 0
    
    try:
        # 最多录制max_record_seconds秒
        for i in range(0, int(RATE / CHUNK * max_record_seconds)):
            data = stream.read(CHUNK, exception_on_overflow=False)
            frames.append(data)
            
            # 计算当前块的音量
            audio_data = np.frombuffer(data, dtype=np.int16)
            volume = np.abs(audio_data).mean()
            
            # 每秒输出一次当前音量（每44次循环约为1秒）
            if i % 44 == 0:
                print(f"当前音量: {volume:.2f}, 阈值: {adjusted_threshold:.2f}, 静音计数: {silent_chunks}/{max_silent_chunks}")
            
            # 检测是否有声音
            if volume > adjusted_threshold:
                silent_chunks = 0  # 重置静音计数
                voice_detected = True
                print(f"检测到声音: {volume:.2f}")
            else:
                silent_chunks += 1
            
            # 如果连续静音超过指定时间，结束录制
            # 注意：如果从未检测到声音，也会在2秒后自动结束
            if silent_chunks >= max_silent_chunks:
                if voice_detected:
                    print(f"检测到 {silence_duration} 秒静音，自动停止录制")
                else:
                    print(f"{silence_duration} 秒内未检测到声音，自动停止录制")
                break
                
    except KeyboardInterrupt:
        print("\n录制被用户中断")
    except Exception as e:
        print(f"录制出错: {str(e)}")
    finally:
        # 停止并关闭流
        print("录制结束...")
        stream.stop_stream()
        stream.close()
        audio.terminate()
        
        # 如果检测到了声音，则保存WAV文件
        if voice_detected:
            print("检测到声音，保存文件中...")
            with wave.open(output_file, 'wb') as wf:
                wf.setnchannels(CHANNELS)
                wf.setsampwidth(audio.get_sample_size(FORMAT))
                wf.setframerate(RATE)
                wf.writeframes(b''.join(frames))
            print(f"音频已保存为 {output_file}")
        else:
            print("未检测到声音，不保存文件")
        
        return voice_detected

def record_and_recognize(output_file='recording.wav', device_index=None, device_keyword="audio", lang_sym="zh", trigger_keys=('space', 'r')):
    """
    录制麦克风音频：点击一次开始录制，再次点击结束录制，并返回语音识别的文字结果
    
    参数:
        output_file: 输出文件名(.wav)
        device_index: 音频设备索引，None表示自动选择
        device_keyword: 设备选择关键词，默认为"audio"
        lang_sym: 语言代码，默认为"zh"（中文）
        trigger_keys: 触发键组合，默认为('space', 'r')，表示同时按下空格键和r键
    
    返回:
        识别的文字结果
    """
    # 如果没有指定设备索引，则自动选择设备
    if device_index is None:
        # 列出所有可用设备
        input_devices = list_audio_devices()
        
        if not input_devices:
            print("无法继续录音，请检查您的麦克风连接或权限设置。")
            return None
        
        # 查找包含关键词的设备
        device_index = find_audio_device(input_devices, device_keyword)
        print(f"将使用设备ID {device_index} 录制音频到文件 {output_file}")
        
    # 音频参数设置
    FORMAT = pyaudio.paInt16  # 16-bit采样
    CHANNELS = 1              # 单声道
    RATE = 44100              # 采样率44.1kHz
    CHUNK = 1024              # 每次读取的帧数
    
    audio = pyaudio.PyAudio()
    
    # 打开音频流
    stream = audio.open(format=FORMAT,
                        channels=CHANNELS,
                        rate=RATE,
                        input=True,
                        input_device_index=device_index,
                        frames_per_buffer=CHUNK)
    
    frames = []
    is_recording = False
    stop_recording = False
    
    # 创建停止录音的事件
    stop_event = threading.Event()
    
    # 跟踪哪些键被按下
    pressed_keys = set()
    
    def key_down_callback(e):
        nonlocal pressed_keys
        pressed_keys.add(e.name)
        
        # 检查是否所有触发键都已经按下
        all_keys_pressed = all(key in pressed_keys for key in trigger_keys)
        
        if all_keys_pressed:
            nonlocal is_recording, stop_recording
            if not is_recording:
                # 开始录制
                is_recording = True
                print(f"\n开始录音...（再次同时按下{'+'.join(trigger_keys)}停止）")
            else:
                # 停止录制
                stop_recording = True
                print("\n停止录音...")
    
    def key_up_callback(e):
        nonlocal pressed_keys
        if e.name in pressed_keys:
            pressed_keys.remove(e.name)
    
    # 注册键盘事件监听
    keyboard.on_press(key_down_callback)
    keyboard.on_release(key_up_callback)
    
    try:
        print(f"同时按下 {'+'.join(trigger_keys)} 开始录音...")
        
        # 等待开始录制
        while not is_recording and not stop_event.is_set():
            time.sleep(0.1)
        
        # 如果用户关闭了程序，则退出
        if stop_event.is_set():
            return None
            
        # 主录音循环
        print("正在录音中...")
        while is_recording and not stop_recording and not stop_event.is_set():
            data = stream.read(CHUNK, exception_on_overflow=False)
            frames.append(data)
            
    except KeyboardInterrupt:
        print("\n录制被用户中断")
    except Exception as e:
        print(f"录制出错: {str(e)}")
    finally:
        # 停止并关闭流
        keyboard.unhook_all()  # 移除键盘事件监听
        stream.stop_stream()
        stream.close()
        audio.terminate()
        
        # 如果有录制到内容，则保存为WAV文件
        if frames:
            print("录制结束，保存文件中...")
            with wave.open(output_file, 'wb') as wf:
                wf.setnchannels(CHANNELS)
                wf.setsampwidth(audio.get_sample_size(FORMAT))
                wf.setframerate(RATE)
                wf.writeframes(b''.join(frames))
            
            print(f"音频已保存为 {output_file}")
            
            # 调用语音识别函数
            try:
                # 导入dolphin_asr模块
                import dolphin_asr
                
                print("正在进行语音识别...")
                # 加载模型并进行识别
                result, _ = dolphin_asr.stream_asr(output_file, lang_sym=lang_sym)
                recognized_text = result.text
                
                print(f"识别结果: {recognized_text}")
                return recognized_text
            except ImportError:
                print("警告: 无法导入dolphin_asr模块，无法进行语音识别")
                return None
            except Exception as e:
                print(f"语音识别出错: {str(e)}")
                return None
        else:
            print("未录制任何内容")
            return None

if __name__ == "__main__":
    # 使用示例
    output_filename = "my_recording.wav"
    
    # 使用新的带语音识别的录音函数，设置结束录制触发键为空格+r
    result_text = record_and_recognize(output_filename, device_index=None, lang_sym="zh", trigger_keys=('space', 'r'))
    
    if result_text:
        print(f"识别的文本: {result_text}")
    else:
        print("未能获取识别结果。")